# -*- coding: utf-8 -*-
"""
Created on Sat Nov 28 22:08:54 2020

@author: Administrator
"""

import numpy as np


def distribution(theta, head, tail):
    """
    coins follow bernoulli ditribution
    """
    return theta**head * (1-theta)**tail
    
    
def em_single(priors, observations):
    """
    em algrithem for single step
    input:
        priors: priors of coins [a, b]
        observations: observations
    """
    # inistial number
    numerator_a = 0
    denominator_a = 0
    
    numerator_b = 0
    denominator_b = 0
    
    # 1 row is 1 case
    for obs in observations:
        #p(x) = sum(ak* p^head * (1-p)^tail)
        # number of heads
        head = obs.sum()
        # number of tails
        tail = len(obs) - head
        # p(z|x,θ) = p(x, z) / p(x) = ak * p^head * (1-p)^tail / p(x)
        Pa = distribution(priors[0], head, tail)  
        Pb = distribution(priors[1], head, tail)
        
        # u_a = p(z=a|x) / p(x)
        # we also call it the respones of a model
        u_a = Pa / (Pa + Pb)
        u_b = 1 - u_a
        
        # partial derivative equals 0 dQ(θ，θ.t)/dpa
        numerator_a += u_a * head
        denominator_a += u_a * (head + tail)
        # partial derivative equals 0 dQ(θ，θ.t)/dpb
        numerator_b += u_b * head
        denominator_b += u_b * (head + tail)
        
    new_pri_a = numerator_a / denominator_a
    new_pri_b = numerator_b / denominator_b

    return [new_pri_a, new_pri_b]          
    
    
def em(priors, observations, tol=1e-6, max_iter=500):
    step = 0
    
    while step < max_iter:
        # iteration
        new_priors = em_single(priors, observations)
        # max change
        det = max(new_priors[0] - priors[0], new_priors[1] - priors[1])
        priors = new_priors
        step += 1
        # if change is small then stop iteration
        if det < tol:
            break
        
    return [new_priors, step]


if __name__ == '__main__':
    # in this case The probability of choosing a is the same as choosing b
    observations = np.array([[1, 0, 0, 0, 1, 1, 0, 1, 0, 1],
                         [1, 1, 1, 1, 0, 1, 1, 1, 1, 1],
                         [1, 0, 1, 1, 1, 1, 1, 0, 1, 1],
                         [1, 0, 1, 0, 0, 0, 1, 1, 0, 0],
                         [0, 1, 1, 1, 0, 1, 1, 1, 0, 1]])
    priors = [0.6, 0.4]
    print (em(priors, observations))
